# GF_OPEN_WEBUI_0_6_Pipelines_Integrations_Dify_Workflow.py
# Crate by GF 2025-04-30 23:30

from typing import List, Union, Generator, Iterator, Optional
from pprint import pprint  # 美化打印工具
import requests, json, warnings

# 如需禁用 SSL 验证警告可取消注释 (生产环境不建议) 
# warnings.filterwarnings('ignore', message='Unverified HTTPS request')

class Pipeline:
    def __init__(self):
        # 管道基础配置
        self.name = "Dify Agent Pipeline"  # 管道名称标识
        self.api_url = "http://dify.hostname/v1/workflows/run"  # Dify 工作流 API 地址 (需替换实际地址) 
        self.api_key = "app-dify-key"      # Dify 应用 API 密钥 (需替换真实密钥) 
        self.api_request_stream = True     # 是否启用流式响应 (与 Dify 服务端配置匹配) 
        self.verify_ssl = True             # SSL 证书验证开关 (生产环境应保持 True) 
        self.debug = False                 # 调试模式开关

    # 服务生命周期管理 --------------------------------------------------
    async def on_startup(self):
        """ 服务启动时触发 (可初始化资源) """
        print(f"on_startup: {__name__}")  # 打印启动日志
        # 可在此处添加异步初始化逻辑, 如建立数据库连接等

    async def on_shutdown(self): 
        """ 服务关闭时触发 (可释放资源) """
        print(f"on_shutdown: {__name__}")  # 打印关闭日志
        # 可在此处添加清理逻辑, 如关闭文件句柄等

    # 数据预处理/后处理 -------------------------------------------------
    async def inlet(self, body: dict, user: Optional[dict] = None) -> dict:
        """ OpenAI 请求前数据预处理 """
        print(f"inlet: {__name__}")
        if self.debug:  # 调试模式下打印原始输入
            print(f"inlet: {__name__} - body:")
            pprint(body)  # 美化打印请求体
            print(f"inlet: {__name__} - user:")
            pprint(user)  # 美化打印用户信息
        # 可在此处修改请求参数, 如添加系统提示词等
        return body  # 返回处理后的请求体

    async def outlet(self, body: dict, user: Optional[dict] = None) -> dict:
        """ OpenAI 响应后数据处理 """
        print(f"outlet: {__name__}")
        if self.debug:  # 调试模式下打印原始输出
            print(f"outlet: {__name__} - body:")
            pprint(body)
            print(f"outlet: {__name__} - user:")
            pprint(user)
        # 可在此处添加响应数据清洗 / 格式化逻辑
        return body  # 返回处理后的响应体

    # 核心处理管道 -----------------------------------------------------
    def pipe(
        self, user_message: str, model_id: str, messages: List[dict], body: dict
    ) -> Union[str, Generator, Iterator]:
        """ Dify 工作流调用主逻辑 """
        print(f"pipe: {__name__}")
        
        # 调试模式打印原始输入
        if self.debug:
            print(f"pipe: {__name__} - received message from user: {user_message}")

        # 设置响应模式 (流式 / 阻塞) 
        response_mode = "streaming" if self.api_request_stream else "blocking"

        # 构建API请求参数
        headers = {
            'Authorization': f'Bearer {self.api_key}',  # API 密钥认证
            'Content-Type': 'application/json'          # JSON 格式请求体
        }
        data = {
            "inputs": {"prompt": user_message},  # 用户输入作为工作流输入
            "response_mode": response_mode,      # 响应模式设置
            "user": body["user"]["email"]        # 从请求体提取用户标识
        }

        # 发送 HTTP POST 请求
        response = requests.post(
            self.api_url,
            headers=headers,
            json=data,
            stream=self.api_request_stream,  # 是否流式接收
            verify=self.verify_ssl           # SSL 验证控制
        )

        # 处理响应
        if response.status_code == 200:
            # 流式处理响应内容
            for line in response.iter_lines():
                if line:  # 过滤空行
                    try:
                        # 解析 SSE 格式数据 (去除 "data: " 前缀) 
                        json_data = json.loads(line.decode('utf-8').replace('data: ', ''))
                        # 提取嵌套的 text 字段 (假设 Dify 返回格式为{"data": {"text": ...}}) 
                        if 'data' in json_data and 'text' in json_data['data']:
                            yield json_data['data']['text']  # 流式返回文本内容
                    except json.JSONDecodeError:
                        print(f"Failed to parse JSON: {line}")
                        # 可添加错误处理逻辑, 如返回错误提示
        else:
            # 返回错误信息 (非200状态码) 
            yield f"Workflow request failed with status code: {response.status_code}"
